"""
AI客服WebSocket消费者
处理医疗问答请求
"""
import json
import logging
from datetime import datetime
import time
from channels.generic.websocket import WebsocketConsumer
from channels.exceptions import StopConsumer
from asgiref.sync import async_to_sync

from .qa_service import qa_service

logger = logging.getLogger(__name__)

class AIConsumer(WebsocketConsumer):
    """AI客服WebSocket消费者"""
    
    def __init__(self, *args, **kwargs):
        super().__init__(*args, **kwargs)
        self.group_name = None
        self.user_id = None
    
    def websocket_connect(self, message):
        """处理WebSocket连接"""
        try:
            # 获取用户ID
            self.group_name = self.scope['url_route']['kwargs'].get('group', 'ai_group')
            self.user_id = self.group_name
            
            # 加入组
            async_to_sync(self.channel_layer.group_add)(
                self.group_name,
                self.channel_name
            )
            
            # 接受连接
            self.accept()
            
            # 发送欢迎消息
            self.send_message("welcome", "您好，我是医疗AI助手，请问有什么可以帮您解答的问题？")
            
            logger.info(f"AI助手连接成功: {self.user_id}")
            
        except Exception as e:
            logger.error(f"AI助手连接错误: {str(e)}")
            raise StopConsumer()
    
    def stream_callback(self, chunk):
        """处理流式回答的回调函数"""
        try:
            if not chunk:
                logger.warning("收到空的流式回答片段")
                return
                
            logger.info(f"发送流式回答片段: {chunk[:20]}...")  # 只记录前20个字符
            
            # 发送增量回答
            message_data = {
                "type": "streaming",
                "message": chunk
            }
            
            message_json = json.dumps(message_data, ensure_ascii=False)
            logger.debug(f"发送WebSocket消息: {message_json[:50]}...")  # 只记录前50个字符
            
            # 确保连接还是打开的
            if self.channel_name:
                self.send(message_json)
            else:
                logger.error("WebSocket连接已关闭，无法发送消息")
        except Exception as e:
            logger.error(f"发送流式回答出错: {str(e)}")
    
    def send_message(self, message_type, message_content):
        """发送消息的通用方法"""
        try:
            message_data = {
                "type": message_type,
                "message": message_content
            }
            
            message_json = json.dumps(message_data, ensure_ascii=False)
            logger.debug(f"发送{message_type}消息: {message_json[:50]}...")  # 只记录前50个字符
            
            self.send(message_json)
            return True
        except Exception as e:
            logger.error(f"发送{message_type}消息出错: {str(e)}")
            return False
    
    def websocket_receive(self, message):
        """处理接收到的消息"""
        try:
            # 获取消息文本
            text_data = message.get('text', '')
            if not text_data:
                logger.warning("收到空消息")
                return
                
            # 解析JSON数据
            data = json.loads(text_data)
            
            # 提取问题
            question = data.get('mes', '')
            if not question:
                logger.warning("收到的消息中没有问题内容")
                return
                
            # 记录问题
            logger.info(f"收到问题 - 用户ID: {self.user_id}, 问题: {question}")
            
            # 发送思考中提示
            self.send_message("typing", "正在思考中...")
            
            # 获取AI回答
            start_time = time.time()
            
            # 确保QA服务已初始化
            if not qa_service.is_initialized:
                logger.info("QA服务尚未初始化，正在初始化...")
                success = qa_service.initialize()
                if not success:
                    logger.error("QA服务初始化失败")
                    self.send_message("error", "AI助手初始化失败，请稍后再试。")
                    return
            
            # 使用流式回答
            logger.info("开始调用流式回答")
            answer = qa_service.answer(question, self.stream_callback)
            logger.info(f"流式回答调用完成，总回答: {answer[:50]}..." if answer else "无回答")
            
            end_time = time.time()
            
            # 记录响应时间
            response_time = round(end_time - start_time, 2)
            logger.info(f"问题回答完成 - 耗时: {response_time}秒")
            
            # 发送完成信号
            self.send(json.dumps({
                "type": "complete",
                "response_time": response_time
            }, ensure_ascii=False))
            
        except Exception as e:
            logger.error(f"处理消息错误: {str(e)}")
            
            # 发送错误回复
            self.send_message("error", "我暂时不知道快让朱林哥哥来教教我吧~~~")
    
    def websocket_disconnect(self, message):
        """处理连接断开"""
        try:
            # 离开组
            async_to_sync(self.channel_layer.group_discard)(
                self.group_name,
                self.channel_name
            )
            
            logger.info(f"AI助手连接断开: {self.user_id}")
            
        except Exception as e:
            logger.error(f"断开连接错误: {str(e)}")
            
        # 停止消费者
        raise StopConsumer() 